import os
import pandas as pd
import time
import json
from typing import List, Dict, Any
from openai import OpenAI
from tqdm import tqdm
import logging
from config import Config
from agent_generator import AgentConfigGenerator
from question_types import QuestionParser, ResponseGenerator, QuestionType

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('questionnaire_system.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class DeepSeekClient:
    def __init__(self, config: Config):
        self.client = OpenAI(
            api_key=config.DEEPSEEK_API_KEY,
            base_url=config.DEEPSEEK_API_BASE
        )
        self.config = config
        
    def generate_response(self, prompt: str, agent_profile: Dict[str, str]) -> str:
        """使用DeepSeek API生成回答"""
        system_message = f"""你是一个问卷调查参与者。请根据你的背景和性格特点来回答问题。
        
背景：{agent_profile['background']}
性格：{agent_profile['personality']}

请以第一人称回答，保持自然真实的语气，体现你的背景和性格特点。"""

        for attempt in range(self.config.MAX_RETRIES):
            try:
                response = self.client.chat.completions.create(
                    model="deepseek-chat",
                    messages=[
                        {"role": "system", "content": system_message},
                        {"role": "user", "content": prompt}
                    ],
                    temperature=0.7,
                    max_tokens=2000,
                    timeout=self.config.REQUEST_TIMEOUT
                )
                return response.choices[0].message.content.strip()
            except Exception as e:
                logger.warning(f"第 {attempt + 1} 次请求失败: {str(e)}")
                if attempt < self.config.MAX_RETRIES - 1:
                    time.sleep(2 ** attempt)  # 指数退避
                else:
                    logger.error(f"所有重试失败: {str(e)}")
                    return f"回答生成失败: {str(e)}"

class QuestionnaireSystem:
    def __init__(self):
        self.config = Config()
        self.deepseek_client = DeepSeekClient(self.config)
        self.agent_generator = AgentConfigGenerator(self.config)
        self.questions = []
        self.parsed_questions = []
        self.results = []
        self.agents = []
        
    def load_questions(self, file_path: str) -> List[Dict[str, Any]]:
        """从文件加载问题列表并解析题型"""
        # try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 按问题分割（假设每个问题之间有空行或明显分隔）
        raw_questions = self._split_questions(content)
        
        parsed_questions = []
        for i, question_text in enumerate(raw_questions, 1):
            if question_text.strip():
                parsed_question = QuestionParser.parse_question(question_text)
                parsed_question["original_text"] = question_text.strip()
                parsed_question["question_number"] = i
                parsed_questions.append(parsed_question)
                
                logger.info(f"解析问题 {i}: {parsed_question['type'].value}")
        
        self.parsed_questions = parsed_questions
        logger.info(f"成功解析 {len(parsed_questions)} 个问题")
        return parsed_questions
            
        # except Exception as e:
        #     logger.error(f"加载问题文件失败: {str(e)}")
        #     raise
    
    def _split_questions(self, content: str) -> List[str]:
        """分割问题文本"""
        # 多种分割方式
        import re
        splits = re.split(r'\n\s*\n|\d+[\.．]|\d+、', content)
        return [q.strip() for q in splits if q.strip()]
    
    def generate_agents_based_on_questionnaire(self, num_agents: int = 8, use_auto_generate: bool = True):
        """生成基于问卷的智能体配置"""
        if use_auto_generate:
            logger.info("使用自动生成的智能体配置...")
            self.agents = self.agent_generator.analyze_questionnaire_and_generate_agents(
                self.dir_path+"intro.txt", num_agents
            )
        else:
            logger.info("使用预设的智能体配置...")
            self.agents = self.config.AGENTS[:num_agents]
        
        # 显示生成的智能体信息
        logger.info("生成的智能体配置：")
        for i, agent in enumerate(self.agents, 1):
            logger.info(f"{i}. {agent['name']} - {agent['background']}")
    
    def create_question_prompt(self, questions: List[Dict[str, Any]]) -> str:
        """创建包含所有问题的提示词"""
        prompt = "请回答以下问卷问题：\n\n"
        
        for i, question in enumerate(questions, 1):
            prompt += f"问题{i} ({question['type'].value}):\n"
            prompt += f"{question['text']}\n"
            
            if question['type'] in [QuestionType.SINGLE_CHOICE, QuestionType.MULTIPLE_CHOICE]:
                for option in question.get('options', []):
                    prompt += f"  {option['key']}. {option['text']}\n"
            elif question['type'] == QuestionType.MATRIX_SCALE:
                prompt += f"  量表: {' → '.join(question.get('scale_labels', []))}\n"
                for j, statement in enumerate(question.get('statements', []), 1):
                    prompt += f"  {j}. {statement}\n"
            
            prompt += "\n"
        
        prompt += "请根据您的背景和性格特点，对每个问题给出真实、一致的回答。"
        prompt += "\n\n"
        prompt += "回答格式为：\n"
        prompt += "只需在回答中选择或填写选项或量表的对应数字或字母即可。\n"
        prompt += "如：“问题1：你喜欢的颜色是？[单选题] A. 红色  B. 蓝色  C. 绿色  D. 黄色”\n"
        prompt += "回答：“A”\n"
        prompt += "如：“问题2：你喜欢的运动项目有？[多选题] A. 篮球  B. 足球  C. 乒乓球  D. 羽毛球”\n"
        prompt += "回答：“A,B”\n"
        prompt += """如：“您希望AI智能体在管理转型中主要解决哪些问题?[矩阵量表题] *
	非常不需要	不需要	一般	需要	非常需要
整合分析数据	○	○	○	○	○
监控风险，自动预警	○	○	○	○	○
模拟、迭代转型方案，辅助决策	○	○	○	○	○"""
        prompt += "回答：“整合分析数据: 非常需要;\n 监控风险，自动预警: 一般;\n 模拟、迭代转型方案，辅助决策: 一般;\n”\n"

        return prompt
    
    def process_agent_responses(self) -> List[Dict[str, Any]]:
        """处理所有智能体的回答"""
        if not self.agents:
            logger.warning("没有智能体配置，使用默认配置")
            self.agents = self.config.AGENTS
        
        results = []
        question_prompt = self.create_question_prompt(self.parsed_questions)
        
        logger.info(f"开始处理 {len(self.agents)} 个智能体的回答...")
        
        for agent in tqdm(self.agents, desc="处理智能体"):
            try:
                logger.info(f"正在处理智能体: {agent['name']}")
                
                # 生成回答
                response = self.deepseek_client.generate_response(question_prompt, agent)
                
                # 解析回答
                answers = self.parse_answers(response, len(self.parsed_questions))
                
                # 构建结果记录
                result_record = {
                    "智能体名称": agent['name'],
                    "背景": agent['background'],
                    "性格": agent['personality'],
                    "完整回答": response
                }
                
                # 添加每个问题的单独回答
                for i, (question, answer) in enumerate(zip(self.parsed_questions, answers), 1):
                    # 为选择题添加选项字段
                    if question['type'] in [QuestionType.SINGLE_CHOICE, QuestionType.MULTIPLE_CHOICE]:
                        result_record[f"问题{i}选择"] = self._extract_choice_from_answer(answer, question)
                    elif question['type'] == QuestionType.MATRIX_SCALE:
                        result_record[f"问题{i}回答"] = answer
                        answer_map = self._extract_choice_from_matrix_answer(answer, question)
                        for j, label in enumerate(question.get('statements', []), 1):
                            result_record[f"问题{i}-{label}"] = answer_map.get(label, "")
                    else:
                        result_record[f"问题{i}回答"] = answer
                
                results.append(result_record)
                logger.info(f"智能体 {agent['name']} 回答完成")
                
                # 避免请求过于频繁
                time.sleep(1)
                
            except Exception as e:
                logger.error(f"处理智能体 {agent['name']} 时出错: {str(e)}")
                # 添加错误记录
                error_record = {
                    "智能体名称": agent['name'],
                    "背景": agent['background'],
                    "性格": agent['personality'],
                    "完整回答": f"处理失败: {str(e)}",
                    "错误": "是"
                }
                results.append(error_record)
        
        self.results = results
        return results
    
    def _extract_choice_from_answer(self, answer: str, question: Dict[str, Any]) -> str:
        """从回答中提取选择项"""
        # 简单的关键词匹配
        options = question.get('options', [])
        selected = []
        answer = answer.split("（")[0]
        for option in options:
            if option['key'] in answer or option['text'] in answer:
                selected.append(option['key'])
        
        return ','.join(selected) if selected else answer
    
    def _extract_choice_from_matrix_answer(self, answer: str, question: Dict[str, Any]) -> str:
        """从回答中提取量表答案"""
        # 简单的关键词匹配
        answers = answer.split(';')
        states = question.get('statements', [])
        mapping_list = []
        idx = 0
        for answer_item in answers:
            if ':' in answer_item:
                key, value = answer_item.split(':')
                key = key.strip()
                value = value.strip()
                mapping_list.append(value)
                assert key in states[idx] or states[idx] in key, f"量表答案 {key} 与问题 {question['text']} 不匹配"
                idx += 1
        
        mapping = {states[i]: mapping_list[i] for i in range(len(mapping_list))}
        
        return mapping
        
    
    def parse_answers(self, response: str, expected_count: int) -> List[str]:
        """解析回答文本，提取每个问题的答案"""
        # 按问题编号分割
        lines = response.split('\n')
        answers = []
        current_answer = ""
        current_question = 0
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            import re
            # 检查是否是新的问题回答开始
            question_match = re.match(r'^\**\s*问题\s*(\d+)\s*\**[：:.．]?\s*(.*)', line)
            if question_match:
                if current_answer and current_question > 0:
                    answers.append(current_answer.strip())
                    current_answer = ""
                
                current_question = int(question_match.group(1))
                rest_of_line = question_match.group(2)
                if rest_of_line:
                    current_answer += rest_of_line + " "
            else:
                current_answer += line + " "
        
        if current_answer:
            answers.append(current_answer.strip())
        
        # 补齐缺失的回答
        while len(answers) < expected_count:
            answers.append("")
        
        # 截断多余的回答
        answers = answers[:expected_count]
        
        return answers
    
    def save_to_excel(self, output_file: str = "问卷结果.xlsx"):
        """保存结果到Excel文件"""
        try:
            if not self.results:
                logger.warning("没有结果数据可保存")
                return
            
            # 创建DataFrame
            df = pd.DataFrame(self.results)
            
            # 重新排列列的顺序
            columns_order = ['智能体名称', '背景', '性格', '完整回答']
            other_columns = [col for col in df.columns if col not in columns_order]
            final_columns = columns_order + other_columns
            
            df = df[final_columns]
            
            # 保存到Excel
            with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
                df.to_excel(writer, sheet_name='回答结果', index=False)
                
                # 添加问题信息表
                question_info = []
                for i, question in enumerate(self.parsed_questions, 1):
                    question_info.append({
                        '问题编号': i,
                        '问题类型': question['type'].value,
                        '问题内容': question['text'][:100] + '...' if len(question['text']) > 100 else question['text'],
                        '选项/量表': str(question.get('options', question.get('scale_labels', '')))
                    })
                
                question_df = pd.DataFrame(question_info)
                question_df.to_excel(writer, sheet_name='问题信息', index=False)
            
            logger.info(f"结果已保存到: {output_file}")
            
            # 同时保存CSV格式作为备份
            csv_file = output_file.replace('.xlsx', '.csv')
            df.to_csv(csv_file, index=False, encoding='utf-8-sig')
            logger.info(f"备份CSV文件已保存到: {csv_file}")
            
            # 保存智能体配置
            agents_file = "生成的智能体配置.json"
            with open(agents_file, 'w', encoding='utf-8') as f:
                json.dump(self.agents, f, ensure_ascii=False, indent=2)
            logger.info(f"智能体配置已保存到: {agents_file}")
            
        except Exception as e:
            logger.error(f"保存文件失败: {str(e)}")
            raise
    
    def run(self, dir_path: str, questions_file: str, output_file: str = "问卷结果.xlsx", 
            num_agents: int = 8, auto_generate_agents: bool = True):
        """运行整个系统"""
        # try:
        logger.info("开始运行LLM批量填问卷系统")
        questions_file = os.path.join(dir_path, questions_file)
        output_file = os.path.join(dir_path, output_file)
        self.dir_path = dir_path
        # 加载问题
        self.load_questions(questions_file)
        
        if not self.parsed_questions:
            logger.error("没有加载到任何问题")
            return
        
        # 生成智能体配置
        self.generate_agents_based_on_questionnaire(num_agents, auto_generate_agents)
        
        # 处理回答
        self.process_agent_responses()
        
        # 保存结果
        self.save_to_excel(output_file)
        
        logger.info("系统运行完成")
        # 加载问题
        self.load_questions(questions_file)
        
        if not self.parsed_questions:
            logger.error("没有加载到任何问题")
            return
        
        # 生成智能体配置
        self.generate_agents_based_on_questionnaire(num_agents, auto_generate_agents)
        
        # 处理回答
        self.process_agent_responses()
        
        # 保存结果
        self.save_to_excel(output_file)
        
        logger.info("系统运行完成")
        
        # except Exception as e:
        #     logger.error(f"系统运行失败: {str(e)}")
        #     raise

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='LLM批量填问卷系统')
    parser.add_argument('--dir', '-d', default='queryB/', 
                       help='问题目录路径')
    parser.add_argument('--questions', '-q', default='questions.txt', 
                       help='问题文件路径')
    parser.add_argument('--output', '-o', default='results.xlsx', 
                       help='输出文件路径')
    parser.add_argument('--agents', '-a', type=int, default=200, 
                       help='智能体数量')
    parser.add_argument('--preset', '-p', action='store_true', 
                       help='使用预设智能体配置（默认使用自动生成）')
    
    args = parser.parse_args()
    
    system = QuestionnaireSystem()
    
    try:
        system.run(
            dir_path=args.dir,
            questions_file=args.questions,
            output_file=args.output,
            num_agents=args.agents,
            auto_generate_agents=not args.preset
        )
        print(f"\n✅ 完成！结果已保存到: {args.dir + args.output}")
    except Exception as e:
        print(f"\n❌ 运行失败: {str(e)}")
        logger.error(f"主程序运行失败: {str(e)}")

if __name__ == "__main__":
    main()